import json
import os

import numpy as np
import torch
import sys
import time
import math
from typing import List, Optional, Sequence, Tuple, Union, Dict

import open3d as o3d
from open3d import geometry
from mmengine.visualization import Visualizer
import matplotlib.pyplot as plt
import colorsys
import cv2 as cv

class OccupancyVisualizer(Visualizer):
    def __init__(self,
                 name: str = 'visualizer',
                 background_color=(0, 0, 0),
                 image: Optional[np.ndarray] = None,
                 vis_backends: Optional[Dict] = None,
                 save_dir: Optional[str] = None,
                 color_map=None):
        super().__init__(
            name=name,
            image=image,
            vis_backends=vis_backends,
            save_dir=save_dir)
        color_map = np.array(color_map)
        color_map = color_map[:, :3]
        self.background_color = background_color
        self.color_map = color_map

        self.flag_pause = True  # 默认暂停，第一帧静止
        self.flag_next = False
        self.flag_exit = False
        self.first_frame = True
        self.view_control = None
        self.o3d_vis = None
        self.current_view = None
        self.initial_view = None
        self.last_update_time = 0
        self.frame_interval = 1.0 / 30.0  # 30fps
        
        # 几何体引用
        self.voxel_grid = None
        self.line_sets = None
        self.mesh_frame = None
        self.car_model_mesh = None
        self.ego_line_set = None
        self.pcd = None

    def escape_callback(self, vis):
        self.flag_exit = True
        return False

    def space_action_callback(self, vis, action, mods):
        if action == 1:  # 按下空格
            if self.flag_pause:
                print('继续播放，按[空格]暂停')
            else:
                print('暂停播放，按[空格]继续')
            self.flag_pause = not self.flag_pause
            self.flag_next = True  # 确保状态改变时触发更新
            # 在状态切换时重新应用渲染选项
            self._apply_render_options()
        return False

    def _clear_o3d_vis(self) -> None:
        """Clear open3d vis."""
        if hasattr(self, 'o3d_vis') and self.o3d_vis is not None:
            # 保存当前视角状态
            if self.view_control is not None:
                self.current_view = self.view_control.convert_to_pinhole_camera_parameters()
            
            # 清除所有几何体
            self.o3d_vis.clear_geometries()
            
            # 恢复视角状态
            if self.current_view is not None:
                self.view_control.convert_from_pinhole_camera_parameters(self.current_view)

    def _initialize_o3d_vis(self):
        """Initialize open3d vis according to frame_cfg."""
        if o3d is None or geometry is None:
            raise ImportError(
                'Please run "pip install open3d" to install open3d first.')
        glfw_key_escape = 256  # Esc
        glfw_key_space = 32  # Space
        o3d_vis = o3d.visualization.VisualizerWithKeyCallback()
        o3d_vis.register_key_callback(glfw_key_escape, self.escape_callback)
        o3d_vis.register_key_action_callback(glfw_key_space,
                                             self.space_action_callback)
        # 增加窗口分辨率到4K
        o3d_vis.create_window(width=3840, height=2160)

        # 设置渲染选项
        render_option = o3d_vis.get_render_option()
        render_option.mesh_show_back_face = True
        render_option.show_coordinate_frame = True
        render_option.background_color = np.asarray(self.background_color)
        render_option.point_size = 4.0
        render_option.line_width = 3.0
        render_option.light_on = True
        render_option.mesh_show_wireframe = False  # 关闭网格线显示
        
        # 使用当前版本支持的渲染选项
        try:
            render_option.mesh_shade_option = o3d.visualization.MeshShadeOption.Color
            render_option.mesh_color_option = o3d.visualization.MeshColorOption.Color
            render_option.light_on = True
            render_option.point_show_normal = False
            render_option.line_width = 10.0  # 增加线宽以提高清晰度
        except AttributeError:
            print("警告：某些渲染选项在当前Open3D版本中不可用")

        # 设置视图控制选项
        view_control = o3d_vis.get_view_control()
        view_control.set_zoom(0.7)
        view_control.set_front([0, 0, 1])
        view_control.set_lookat([0, 0, 40])
        view_control.set_up([0, 1, 0])

        # 设置初始视角
        if hasattr(self, 'current_view') and self.current_view is not None:
            view_control.convert_from_pinhole_camera_parameters(self.current_view)

        self.view_control = view_control
        return o3d_vis

    def _apply_render_options(self):
        """应用渲染选项"""
        if hasattr(self, 'o3d_vis') and self.o3d_vis is not None:
            render_option = self.o3d_vis.get_render_option()
            render_option.mesh_show_back_face = True
            render_option.show_coordinate_frame = True
            render_option.background_color = np.asarray(self.background_color)
            render_option.point_size = 4.0
            render_option.line_width = 10.0  # 保持较粗的线条宽度
            render_option.light_on = True
            # 添加网格线相关设置
            render_option.mesh_show_wireframe = False  # 关闭网格线显示

    def flow_to_color(self, vx, vy, max_magnitude=None):
        magnitude = np.sqrt(vx ** 2 + vy ** 2)
        angle = np.arctan2(vy, vx)

        # 将方向映射到色相 (0 到 1)
        hue = (angle + np.pi) / (2 * np.pi)

        # 将大小映射到亮度或饱和度
        if max_magnitude is None:
            max_magnitude = np.max(magnitude)
        if max_magnitude == 0:
            saturation = np.zeros_like(magnitude)
        else:
            saturation = np.clip(magnitude / max_magnitude + 1e-3, 0, 1)
        value = np.ones_like(saturation)

        # 将 HSV 转换为 RGB
        hsv = np.stack((hue, saturation, value), axis=-1)
        rgb = np.apply_along_axis(lambda x: colorsys.hsv_to_rgb(*x), -1, hsv)
        rgb = np.nan_to_num(rgb, nan=0.0, posinf=1.0, neginf=0.0)

        # 将 RGB 转换为 0-255 范围的整数
        rgb = (rgb * 255).astype(np.uint8)

        return rgb

    def _voxel2points(self,
                      voxel,
                      voxel_flow,
                      voxelSize=(0.4, 0.4, 0.4),
                      range=[-40.0, -40.0, -1.0, 40.0, 40.0, 5.4],
                      ignore_labels=[17, 255]):
        if isinstance(voxel, np.ndarray): voxel = torch.from_numpy(voxel)
        mask = torch.zeros_like(voxel, dtype=torch.bool)
        for ignore_label in ignore_labels:
            mask = torch.logical_or(voxel == ignore_label, mask)
        mask = torch.logical_not(mask)
        occIdx = torch.where(mask)
        points = torch.cat((occIdx[0][:, None] * voxelSize[0] + voxelSize[0] / 2 + range[0], \
                            occIdx[1][:, None] * voxelSize[1] + voxelSize[1] / 2 + range[1], \
                            occIdx[2][:, None] * voxelSize[2] + voxelSize[2] / 2 + range[2]), dim=1)
        if voxel_flow is not None:
            flow = voxel_flow[occIdx]
        else:
            flow = None
        return points, voxel[occIdx], flow

    def _voxel_profile(self,
                       voxel,
                       voxel_size=(0.4, 0.4, 0.4)
                       ):
        centers = torch.cat((voxel[:, :2], voxel[:, 2][:, None] - voxel_size[2] / 2), dim=1)
        # centers = voxel
        wlh = torch.cat((torch.tensor(voxel_size[0]).repeat(centers.shape[0])[:, None],
                         torch.tensor(voxel_size[1]).repeat(centers.shape[0])[:, None],
                         torch.tensor(voxel_size[2]).repeat(centers.shape[0])[:, None]), dim=1)
        yaw = torch.full_like(centers[:, 0:1], 0)
        return torch.cat((centers, wlh, yaw), dim=1)

    def _generate_the_ego_car(self):
        # ego_color black

        ego_range = [-1, -1, 0, 3, 1, 1.5]
        ego_voxel_size = [0.1, 0.1, 0.1]
        ego_xdim = int((ego_range[3] - ego_range[0]) / ego_voxel_size[0])
        ego_ydim = int((ego_range[4] - ego_range[1]) / ego_voxel_size[1])
        ego_zdim = int((ego_range[5] - ego_range[2]) / ego_voxel_size[2])
        temp_x = np.arange(ego_xdim)
        temp_y = np.arange(ego_ydim)
        temp_z = np.arange(ego_zdim)
        ego_xyz = np.stack(np.meshgrid(temp_y, temp_x, temp_z), axis=-1).reshape(-1, 3)
        ego_point_x = (ego_xyz[:, 0:1] + 0.5) / ego_xdim * (ego_range[3] - ego_range[0]) + ego_range[0]
        ego_point_y = (ego_xyz[:, 1:2] + 0.5) / ego_ydim * (ego_range[4] - ego_range[1]) + ego_range[1]
        ego_point_z = (ego_xyz[:, 2:3] + 0.5) / ego_zdim * (ego_range[5] - ego_range[2]) + ego_range[2]
        ego_point_xyz = np.concatenate((ego_point_y, ego_point_x, ego_point_z), axis=-1)
        # ego_rgb based on height, it generate rainbow color, the top is red, the bottom is blue
        ego_point_rgb = np.concatenate((ego_point_z, np.zeros_like(ego_point_z), 1 - ego_point_z), axis=-1)
        ego_point_rgb = ego_point_rgb * 255

        # ego_pint_rgb = np.zeros((ego_point_xyz.shape[0], 3))
        # ego_pint_rgb[:, 0] = (ego_point_xyz[:, 2] - ego_range[2]) / (ego_range[5] - ego_range[2])
        # ego_pint_rgb[:, 1] = 1 - ego_pint_rgb[:, 0]
        # ego_point_rgb = ego_pint_rgb * 255

        ego_point_xyzrgb = np.concatenate((ego_point_xyz, ego_point_rgb), axis=-1)

        return ego_point_xyzrgb

    def _my_compute_box_3d(self, center, size, heading_angle):
        h, w, l = size[:, 2], size[:, 0], size[:, 1]
        heading_angle = -heading_angle - math.pi / 2
        center[:, 2] = center[:, 2] + h / 2
        # R = rotz(1 * heading_angle)
        l, w, h = (l / 2).unsqueeze(1), (w / 2).unsqueeze(1), (h / 2).unsqueeze(1)
        x_corners = torch.cat([-l, l, l, -l, -l, l, l, -l], dim=1)[..., None]
        y_corners = torch.cat([w, w, -w, -w, w, w, -w, -w], dim=1)[..., None]
        z_corners = torch.cat([h, h, h, h, -h, -h, -h, -h], dim=1)[..., None]
        # corners_3d = R @ torch.vstack([x_corners, y_corners, z_corners])
        corners_3d = torch.cat([x_corners, y_corners, z_corners], dim=2)
        corners_3d[..., 0] += center[:, 0:1]
        corners_3d[..., 1] += center[:, 1:2]
        corners_3d[..., 2] += center[:, 2:3]
        return corners_3d

    def voxellizer_points(self,
                          points,
                          ego_points,
                          bbox_corners,
                          linesets,
                          voxel_size=0.4,
                          frame_cfg: dict = dict(size=1, origin=[0, 0, 0]),
                          vis_mode='replace',
                          points_size=4,
                          mode='xyzrgb',
                          points_color: Tuple[float] = (0.8, 0.8, 0.8),
                          show_color=True,
                          car_model_mesh=None,
                          line_width=0.0,
                          ):
        if not hasattr(self, 'o3d_vis'):
            self.o3d_vis = self._initialize_o3d_vis()
        else:
            # 清除所有几何体
            self.o3d_vis.clear_geometries()

        render_option = self.o3d_vis.get_render_option()
        if render_option is not None:
            render_option.point_size = points_size
            render_option.background_color = np.asarray(self.background_color)
            render_option.line_width = line_width  # 设置线宽

        points = points.copy()
        pcd = geometry.PointCloud()
        if mode == 'xyz':
            pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            points_colors = np.tile(
                np.array(points_color), (points.shape[0], 1))
        elif mode == 'xyzrgb':
            if ego_points is not None:
                points = np.concatenate([points, ego_points], axis=0)
            pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            points_colors = points[:, 3:6]
            points_colors = points_colors[:, [2, 1, 0]]
            if not ((points_colors >= 0.0) & (points_colors <= 1.0)).all():
                points_colors /= 255.0
        else:
            raise NotImplementedError
        pcd.colors = o3d.utility.Vector3dVector(points_colors)
        
        if show_color:
            # 首先添加体素网格
            voxelGrid = o3d.geometry.VoxelGrid.create_from_point_cloud(pcd, voxel_size=voxel_size)
            self.o3d_vis.add_geometry(voxelGrid)
            
            # 然后添加线框（如果线宽大于0）
            if line_width > 0:
                line_sets = o3d.geometry.LineSet()
                line_sets.points = o3d.open3d.utility.Vector3dVector(bbox_corners.reshape((-1, 3)))
                line_sets.lines = o3d.open3d.utility.Vector2iVector(linesets.reshape((-1, 2)))
                line_sets.paint_uniform_color((0, 0, 0))
                self.o3d_vis.add_geometry(line_sets)
            
            # 最后添加点云
            self.o3d_vis.add_geometry(pcd)
            self.pcd = pcd
            self.points_colors = points_colors
        
        # 添加坐标轴
        mesh_frame = geometry.TriangleMesh.create_coordinate_frame(**frame_cfg)
        self.o3d_vis.add_geometry(mesh_frame)
        
        # 添加车辆模型
        if car_model_mesh is not None:
            self.o3d_vis.add_geometry(car_model_mesh)

    def vis_occ(self,
                occ_semantics,
                occ_flow=None,
                ignore_labels=None,
                voxelSize=(0.4, 0.4, 0.4),
                range=[-40.0, -40.0, -1.0, 40.0, 40.0, 5.4],
                save_path=None,
                wait_time=1,
                view_json=None,
                car_model_mesh=None,
                line_width=0.0,
                ):
        # 初始化可视化器
        if not hasattr(self, 'o3d_vis') or self.o3d_vis is None:
            self.o3d_vis = self._initialize_o3d_vis()
            self.first_frame = True
            # 只在第一帧时设置初始视角
            if view_json is not None:
                self.initial_view = view_json
                self.current_view = view_json
                self.view_control.convert_from_pinhole_camera_parameters(view_json)
        else:
            # 保存当前视角状态
            if self.view_control is not None:
                self.current_view = self.view_control.convert_to_pinhole_camera_parameters()
            # 清除所有几何体
            self.o3d_vis.clear_geometries()

        # 处理点云数据
        points, occ_voxel, occ_flow = self._voxel2points(occ_semantics, occ_flow, ignore_labels=ignore_labels, voxelSize=voxelSize, range=range)
        points = points.numpy()
        occ_voxel = occ_voxel.numpy()
        if occ_flow is None:
            pts_color = self.color_map[occ_voxel.astype(int) % len(self.color_map)]
        else:
            vx = occ_flow[..., 0]
            vy = occ_flow[..., 1]
            pts_color = self.flow_to_color(vx, vy)
        seg_color = np.concatenate([points[:, :3], pts_color], axis=1)
        
        # 处理车辆模型
        if car_model_mesh is None:
            ego_points = self._generate_the_ego_car()
        else:
            ego_points = None

        # 处理边界框
        bboxes = self._voxel_profile(torch.tensor(points), voxel_size=voxelSize)
        bboxes_corners = self._my_compute_box_3d(bboxes[:, 0:3], bboxes[:, 3:6], bboxes[:, 6:7])
        bases_ = torch.arange(0, bboxes_corners.shape[0] * 8, 8)
        edges = torch.tensor([[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6], [6, 7], [7, 4], [0, 4], [1, 5], [2, 6],
                              [3, 7]])
        edges = edges.reshape((1, 12, 2)).repeat(bboxes_corners.shape[0], 1, 1)
        edges = edges + bases_[:, None, None]

        # 添加几何体
        self.voxellizer_points(
            points=seg_color,
            voxel_size=voxelSize[0],
            ego_points=ego_points,
            bbox_corners=bboxes_corners.numpy(),
            linesets=edges.numpy(),
            mode='xyzrgb',
            vis_mode='add',
            show_color=True,
            car_model_mesh=car_model_mesh,
            line_width=line_width,
        )

        # 恢复视角
        if self.first_frame and view_json is not None:
            # 第一帧使用初始视角
            self.view_control.convert_from_pinhole_camera_parameters(view_json)
        elif self.current_view is not None:
            # 其他帧使用保存的视角
            self.view_control.convert_from_pinhole_camera_parameters(self.current_view)

        # 保存图像
        if save_path is not None:
            try:
                os.makedirs(os.path.dirname(save_path), exist_ok=True)
                # 使用Open3D的capture_screen_image方法获取图像
                self.o3d_vis.capture_screen_image(save_path + '.png')
                # 不打印保存信息，让外层统一处理
            except Exception as e:
                print(f"\n保存图像失败: {e}")

        # 处理事件循环
        if self.first_frame:
            self.first_frame = False
            self.show(wait_time=wait_time)
        else:
            self.o3d_vis.poll_events()
            self.o3d_vis.update_renderer()

    def show(self,
             save_path: Optional[str] = None,
             wait_time: int = -1) -> None:

        if hasattr(self, 'o3d_vis') and self.o3d_vis is not None:
            while not self.flag_exit:
                current_time = time.time()
                elapsed = current_time - self.last_update_time

                # 处理事件
                if not self.o3d_vis.poll_events():
                    self.flag_exit = True
                    break

                # 非暂停状态下更新帧
                if not self.flag_pause and elapsed >= self.frame_interval:
                    if hasattr(self, 'on_next_frame'):
                        try:
                            # 保存当前视角
                            if self.view_control is not None:
                                self.current_view = self.view_control.convert_to_pinhole_camera_parameters()
                            
                            # 清除所有几何体
                            self.o3d_vis.clear_geometries()
                            
                            # 更新场景
                            self.on_next_frame()
                            
                            # 恢复视角
                            if self.current_view is not None:
                                self.view_control.convert_from_pinhole_camera_parameters(self.current_view)
                            
                            self.flag_next = True
                            self.last_update_time = current_time
                        except StopIteration:
                            # 播放结束，重置状态
                            print("播放结束")
                            self.flag_pause = False  # 不暂停，继续播放新场景
                            self.flag_next = True   # 触发一次更新
                            # 重置状态，但不关闭窗口
                            self.flag_exit = False
                            self.first_frame = True
                            # 重置所有几何体引用
                            self.voxel_grid = None
                            self.line_sets = None
                            self.mesh_frame = None
                            self.car_model_mesh = None
                            self.ego_line_set = None
                            self.pcd = None
                            # 清理几何体
                            self._clear_o3d_vis()
                            # 重新抛出异常，让外层处理
                            raise
                
                # 更新渲染
                if self.flag_next:
                    self.o3d_vis.update_renderer()
                    self.flag_next = False
                
                # 等待时间
                if wait_time != -1:
                    time.sleep(wait_time)
                else:
                    # 控制帧率
                    time.sleep(max(0, self.frame_interval - elapsed))
            
            # 退出时清理
            if self.flag_exit:
                self.o3d_vis.destroy_window()
                self.o3d_vis.close()
                self._clear_o3d_vis()
                sys.exit(0)
